#include "../SeanNetworkLib/SeanNetwork.h"
#include <iostream>
#include <thread>
#include <atomic>
#include <string>
#include <map>
#include <mutex>
#include <fstream>
#include <sstream>

// 消息ID定义
enum MessageType {
    MSG_HEARTBEAT = 0,
    MSG_TEXT = 1,
    MSG_FILE_REQUEST = 2,
    MSG_FILE_RESPONSE = 3
};

// 模拟文件系统
class FileStore {
private:
    std::map<std::string, std::vector<uint8_t>> files;
    std::mutex filesMutex;

public:
    FileStore() {
        // 预先添加一些示例文件
        AddFile("hello.txt", "Hello, this is a sample text file for testing!");
        AddFile("data.bin", std::string(1024, 'X'));  // 1KB的二进制数据
        AddFile("large.dat", std::string(1024 * 1024, 'Y'));  // 1MB的数据
    }

    void AddFile(const std::string& name, const std::string& content) {
        std::lock_guard<std::mutex> lock(filesMutex);
        std::vector<uint8_t> data(content.begin(), content.end());
        files[name] = data;
    }

    bool GetFile(const std::string& name, std::vector<uint8_t>& outData) {
        std::lock_guard<std::mutex> lock(filesMutex);
        auto it = files.find(name);
        if (it != files.end()) {
            outData = it->second;
            return true;
        }
        return false;
    }

    std::vector<std::string> ListFiles() {
        std::lock_guard<std::mutex> lock(filesMutex);
        std::vector<std::string> result;
        for (const auto& pair : files) {
            result.push_back(pair.first + " (" + std::to_string(pair.second.size()) + " bytes)");
        }
        return result;
    }
};

// 简单的客户端会话管理
class ClientManager {
private:
    std::map<uint64_t, std::string> clients;
    std::mutex clientsMutex;

public:
    void AddClient(uint64_t clientId, const std::string& info) {
        std::lock_guard<std::mutex> lock(clientsMutex);
        clients[clientId] = info;
        std::cout << "Client connected: " << clientId << " " << info << std::endl;
    }

    void RemoveClient(uint64_t clientId) {
        std::lock_guard<std::mutex> lock(clientsMutex);
        auto it = clients.find(clientId);
        if (it != clients.end()) {
            std::cout << "Client disconnected: " << clientId << " " << it->second << std::endl;
            clients.erase(it);
        }
    }

    size_t GetClientCount() {
        std::lock_guard<std::mutex> lock(clientsMutex);
        return clients.size();
    }

    std::vector<uint64_t> GetAllClientIds() {
        std::lock_guard<std::mutex> lock(clientsMutex);
        std::vector<uint64_t> result;
        for (const auto& pair : clients) {
            result.push_back(pair.first);
        }
        return result;
    }
};

// 服务器统计信息结构体
struct ServerStats {
    size_t ActiveSessions;       // 活跃会话数
    size_t TotalSessions;        // 会话总数
    uint64_t MessagesReceived;   // 接收消息数
    uint64_t MessagesSent;       // 发送消息数
    uint64_t TotalBytesReceived; // 总接收字节数
    uint64_t TotalBytesSent;     // 总发送字节数
    size_t KcpSessions;          // KCP会话数
    size_t TcpSessions;          // TCP会话数
    
    ServerStats() 
        : ActiveSessions(0), TotalSessions(0), MessagesReceived(0), MessagesSent(0),
          TotalBytesReceived(0), TotalBytesSent(0), KcpSessions(0), TcpSessions(0) {}
};

// 统计信息输出
void PrintServerStats(ServerStats stats) {
    std::cout << "==== Server Stats ====" << std::endl;
    std::cout << "Active Sessions: " << stats.ActiveSessions << std::endl;
    std::cout << "Total Sessions: " << stats.TotalSessions << std::endl;
    std::cout << "Messages Received: " << stats.MessagesReceived << std::endl;
    std::cout << "Messages Sent: " << stats.MessagesSent << std::endl;
    std::cout << "Total Bytes Received: " << stats.TotalBytesReceived << std::endl;
    std::cout << "Total Bytes Sent: " << stats.TotalBytesSent << std::endl;
    std::cout << "KCP Sessions: " << stats.KcpSessions << std::endl;
    std::cout << "TCP Sessions: " << stats.TcpSessions << std::endl;
    std::cout << "======================" << std::endl;
}

int main(int argc, char* argv[]) {
    std::cout << "NBNetwork Server Test Application" << std::endl;
    std::cout << "=================================" << std::endl;

    // 解析命令行参数
    uint16_t serverPort = 19000;
    if (argc >= 2) {
        serverPort = static_cast<uint16_t>(std::stoi(argv[1]));
    }

    // 创建服务器配置
    NBNetwork::ServerConfig config;
    config.KcpMaxPacketLoss = 3;      // 最大连续丢包数
    config.KcpMaxLatency = 500;       // 最大延迟(ms)
    config.WorkerThreadCount = 4;     // 工作线程数
    config.MaxConnections = 1000;    // 最大连接数
    config.SessionTimeoutMs = 60000;    // 会话超时(ms)
    config.TcpPort = serverPort;     // TCP端口
    config.KcpPort = serverPort;     // KCP端口

    // 创建服务器实例
    auto server = NBNetwork::Net::CreateServer(config);
    if (!server) {
        std::cerr << "Failed to create server" << std::endl;
        return 1;
    }

    // 创建文件存储和客户端管理
    FileStore fileStore;
    ClientManager clientManager;

    // 注册客户端连接回调
    server->SetSessionCallback([&clientManager](NBNetwork::ClientSessionId clientId, NBNetwork::ConnectionState state, const std::string& reason) {
        if (state == NBNetwork::ConnectionState::Connected) {
            clientManager.AddClient(clientId, reason);
        } else if (state == NBNetwork::ConnectionState::Disconnected) {
            clientManager.RemoveClient(clientId);
        }
    });

    // 注册消息处理回调
    server->RegisterMessageHandler(MSG_TEXT, [&server, &clientManager](uint64_t clientId, std::shared_ptr<NBNetwork::NetworkMessage> message) {
        // 解析文本消息
        std::string text = NBNetwork::Net::GetMessageAsString(message);
        std::cout << "Received text from client " << clientId << ": " << text << std::endl;
        
        // 回复确认消息
        auto response = server->CreateMessage(MSG_TEXT, NBNetwork::MessagePriority::Normal);
        std::string responseText = "Server received: " + text;
        response->SetBody(reinterpret_cast<const uint8_t*>(responseText.data()), responseText.size());
        
        // 将消息发送回客户端
        server->SendMsg(clientId, response);
        
        // 如果是广播命令，则发送给所有客户端
        if (text.find("broadcast:") == 0) {
            std::string broadcastText = "BROADCAST: " + text.substr(10);
            auto broadcastMsg = server->CreateMessage(MSG_TEXT, NBNetwork::MessagePriority::Normal);
            broadcastMsg->SetBody(reinterpret_cast<const uint8_t*>(broadcastText.data()), broadcastText.size());
            
            // 获取所有客户端ID并广播
            auto allClients = clientManager.GetAllClientIds();
            for (uint64_t id : allClients) {
                if (id != clientId) { // 排除发送者
                    server->SendMsg(id, broadcastMsg);
                }
            }
        }
    });

    // 处理文件请求
    server->RegisterMessageHandler(MSG_FILE_REQUEST, [&server, &fileStore](uint64_t clientId, std::shared_ptr<NBNetwork::NetworkMessage> message) {
        // 获取请求的文件名
        std::string filename = NBNetwork::Net::GetMessageAsString(message);
        std::cout << "File request from client " << clientId << ": " << filename << std::endl;
        
        // 从文件存储获取文件内容
        std::vector<uint8_t> fileData;
        if (fileStore.GetFile(filename, fileData)) {
            // 创建文件响应消息
            auto response = server->CreateMessage(MSG_FILE_RESPONSE, NBNetwork::MessagePriority::Low);
            response->SetBody(fileData.data(), fileData.size());
            
            // 发送给客户端
            if (server->SendMsg(clientId, response)) {
                std::cout << "File '" << filename << "' sent to client " << clientId 
                         << " (" << fileData.size() << " bytes)" << std::endl;
            } else {
                std::cout << "Failed to send file to client " << clientId << std::endl;
            }
        } else {
            // 文件不存在，发送错误消息
            std::string errorMsg = "File not found: " + filename;
            auto errorResponse = server->CreateMessage(MSG_TEXT, NBNetwork::MessagePriority::High);
            errorResponse->SetBody(reinterpret_cast<const uint8_t*>(errorMsg.data()), errorMsg.size());
            server->SendMsg(clientId, errorResponse);
            std::cout << "File '" << filename << "' not found for client " << clientId << std::endl;
        }
    });

    // 启动服务器
    if (!server->Start()) {
        std::cerr << "Failed to start server on port " << serverPort << std::endl;
        return 1;
    }

    std::cout << "Server started on port " << serverPort << std::endl;
    std::cout << "Type 'help' for available commands" << std::endl;

    std::string line;
    bool running = true;
    
    // 主循环
    while (running && std::getline(std::cin, line)) {
        // 跳过空行
        if (line.empty()) {
            continue;
        }

        // 解析命令
        std::string command = line;
        std::string param;

        size_t spacePos = line.find(' ');
        if (spacePos != std::string::npos) {
            command = line.substr(0, spacePos);
            param = line.substr(spacePos + 1);
        }

        // 执行命令
        if (command == "quit" || command == "exit") {
            std::cout << "Shutting down server..." << std::endl;
            running = false;
        }
        else if (command == "help") {
            std::cout << "Available commands:" << std::endl;
            std::cout << "  help               - Show this help" << std::endl;
            std::cout << "  stats              - Show server statistics" << std::endl;
            std::cout << "  clients            - List connected clients" << std::endl;
            std::cout << "  files              - List available files" << std::endl;
            std::cout << "  broadcast <msg>    - Broadcast message to all clients" << std::endl;
            std::cout << "  quit               - Shutdown server and exit" << std::endl;
        }
        else if (command == "stats") {
            // 获取当前连接数和服务器运行信息
            ServerStats stats;
            stats.ActiveSessions = server->GetConnectionCount();
            stats.TotalSessions = stats.ActiveSessions;  // 简化处理，实际应累计所有会话
            
            PrintServerStats(stats);
        }
        else if (command == "clients") {
            auto clientIds = clientManager.GetAllClientIds();
            std::cout << "Connected clients (" << clientIds.size() << "):" << std::endl;
            for (uint64_t id : clientIds) {
                std::cout << "  Client ID: " << id << std::endl;
            }
        }
        else if (command == "files") {
            auto files = fileStore.ListFiles();
            std::cout << "Available files (" << files.size() << "):" << std::endl;
            for (const auto& file : files) {
                std::cout << "  " << file << std::endl;
            }
        }
        else if (command == "broadcast" && !param.empty()) {
            auto clientIds = clientManager.GetAllClientIds();
            if (clientIds.empty()) {
                std::cout << "No clients connected" << std::endl;
                continue;
            }
            
            // 创建广播消息
            std::string broadcastText = "SERVER: " + param;
            auto message = server->CreateMessage(MSG_TEXT, NBNetwork::MessagePriority::Normal);
            message->SetBody(reinterpret_cast<const uint8_t*>(broadcastText.data()), broadcastText.size());
            
            // 向所有客户端发送消息
            size_t successCount = 0;
            for (uint64_t id : clientIds) {
                if (server->SendMsg(id, message)) {
                    successCount++;
                }
            }
            
            std::cout << "Broadcast message sent to " << successCount << "/" << clientIds.size() << " clients" << std::endl;
        }
        else {
            std::cout << "Unknown command. Type 'help' for available commands" << std::endl;
        }
    }

    // 停止服务器
    std::cout << "Stopping server..." << std::endl;
    server->Stop();
    std::cout << "Server stopped" << std::endl;
    
    return 0;
}